package framework.utils;

import lombok.Getter;
import lombok.SneakyThrows;

import java.io.DataInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Base64.Decoder;
import java.util.Base64.Encoder;

public class RSAUtil {

    /**
     * KEY SIZE 必要时在使用前修改
     */
    public static int KEY_SIZE = 2048;

    /**
     * 获取公钥
     *
     * @param filename
     * @return
     * @throws Exception
     */
    public PublicKey getPublicKey(String filename) throws Exception {
        try (InputStream resourceAsStream = this.getClass().getClassLoader().getResourceAsStream(filename)) {
            try (DataInputStream dis = new DataInputStream(resourceAsStream)) {
                byte[] keyBytes = new byte[resourceAsStream.available()];
                dis.readFully(keyBytes);
                X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
                KeyFactory kf = KeyFactory.getInstance("RSA");
                return kf.generatePublic(spec);
            }
        }
    }

    /**
     * 获取密钥
     *
     * @param filename
     * @return
     * @throws Exception
     */
    public PrivateKey getPrivateKey(String filename) throws Exception {
        try (InputStream resourceAsStream = this.getClass().getClassLoader().getResourceAsStream(filename)) {
            try (DataInputStream dis = new DataInputStream(resourceAsStream)) {
                byte[] keyBytes = new byte[resourceAsStream.available()];
                dis.readFully(keyBytes);
                PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
                KeyFactory kf = KeyFactory.getInstance("RSA");
                return kf.generatePrivate(spec);
            }
        }
    }

    /**
     * 获取公钥
     *
     * @param filename
     * @return
     * @throws Exception
     */
    public PublicKey getPublicKeyFromBase64(String filename) throws Exception {
        try (InputStream resourceAsStream = this.getClass().getClassLoader().getResourceAsStream(filename)) {
            try (DataInputStream dis = new DataInputStream(resourceAsStream)) {
                byte[] keyBytes = new byte[resourceAsStream.available()];
                dis.readFully(keyBytes);

                String base64 = new String(keyBytes);
                keyBytes = Base64.getDecoder().decode(base64);

                X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
                KeyFactory kf = KeyFactory.getInstance("RSA");
                return kf.generatePublic(spec);
            }
        }
    }

    /**
     * 获取密钥
     *
     * @param filename
     * @return
     * @throws Exception
     */
    public PrivateKey getPrivateKeyFromBase64(String filename) throws Exception {
        try (InputStream resourceAsStream = this.getClass().getClassLoader().getResourceAsStream(filename)) {
            try (DataInputStream dis = new DataInputStream(resourceAsStream)) {
                byte[] keyBytes = new byte[resourceAsStream.available()];
                dis.readFully(keyBytes);

                String base64 = new String(keyBytes);
                keyBytes = Base64.getDecoder().decode(base64);

                PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
                KeyFactory kf = KeyFactory.getInstance("RSA");
                return kf.generatePrivate(spec);
            }
        }
    }

    /**
     * 获取公钥
     *
     * @param publicKey
     * @return
     * @throws Exception
     */
    public PublicKey getPublicKey(byte[] publicKey) throws Exception {
        X509EncodedKeySpec spec = new X509EncodedKeySpec(publicKey);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePublic(spec);
    }

    /**
     * 获取密钥
     *
     * @param privateKey
     * @return
     * @throws Exception
     */
    public PrivateKey getPrivateKey(byte[] privateKey) throws Exception {
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(privateKey);
        KeyFactory kf = KeyFactory.getInstance("RSA");
        return kf.generatePrivate(spec);
    }

    /**
     * 生存rsa公钥和密钥
     *
     * @param publicKeyFilename
     * @param privateKeyFilename
     * @param password
     * @throws IOException
     * @throws NoSuchAlgorithmException
     */
    public void generateKey(String publicKeyFilename, String privateKeyFilename, String password) throws IOException, NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        SecureRandom secureRandom = new SecureRandom(password.getBytes());
        keyPairGenerator.initialize(KEY_SIZE, secureRandom);
        KeyPair keyPair = keyPairGenerator.genKeyPair();
        byte[] publicKeyBytes = keyPair.getPublic().getEncoded();
        FileOutputStream fos = new FileOutputStream(publicKeyFilename);
        fos.write(publicKeyBytes);
        fos.close();
        byte[] privateKeyBytes = keyPair.getPrivate().getEncoded();
        fos = new FileOutputStream(privateKeyFilename);
        fos.write(privateKeyBytes);
        fos.close();
    }

    /**
     * 生存rsa公钥和密钥
     *
     * @param publicKeyFilename
     * @param privateKeyFilename
     * @param password
     * @throws IOException
     * @throws NoSuchAlgorithmException
     */
    public void generateKeyBase64(String publicKeyFilename, String privateKeyFilename, String password) throws IOException, NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        SecureRandom secureRandom = new SecureRandom(password.getBytes());
        keyPairGenerator.initialize(KEY_SIZE, secureRandom);
        KeyPair keyPair = keyPairGenerator.genKeyPair();
        byte[] publicKeyBytes = keyPair.getPublic().getEncoded();
        byte[] publicBase64Bytes = Base64.getEncoder().encode(publicKeyBytes);
        try (FileOutputStream fos = new FileOutputStream(publicKeyFilename)) {
            fos.write(publicBase64Bytes);
        }
        byte[] privateKeyBytes = keyPair.getPrivate().getEncoded();
        byte[] privateBase64Bytes = Base64.getEncoder().encode(privateKeyBytes);
        try (FileOutputStream fos = new FileOutputStream(privateKeyFilename)) {
            fos.write(privateBase64Bytes);
        }
    }

    /**
     * 生存rsa公钥
     *
     * @param password
     * @throws IOException
     * @throws NoSuchAlgorithmException
     */
    public static byte[] generatePublicKey(String password) throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        SecureRandom secureRandom = new SecureRandom(password.getBytes());
        keyPairGenerator.initialize(KEY_SIZE, secureRandom);
        KeyPair keyPair = keyPairGenerator.genKeyPair();
        return keyPair.getPublic().getEncoded();
    }

    /**
     * 生存rsa
     *
     * @param password
     * @throws IOException
     * @throws NoSuchAlgorithmException
     */
    public static KeyPair generateKey(String password) throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        SecureRandom secureRandom = new SecureRandom(password.getBytes());
        keyPairGenerator.initialize(KEY_SIZE, secureRandom);
        KeyPair keyPair = keyPairGenerator.genKeyPair();
        return keyPair;
    }

    public static String toBase64(byte[] b) {
        Encoder encoder = Base64.getEncoder();
        String encode = encoder.encodeToString(b);
        return encode;
    }

    public static final byte[] toBytes(String s) {
        Decoder decoder = Base64.getDecoder();
        byte[] buffer = decoder.decode(s);
        return buffer;
    }

    @SneakyThrows
    public static final RSAItem generateNew(int length) {
        // private
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(length);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        PrivateKey privateKey = keyPair.getPrivate();
        //
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        RSAPrivateKeySpec privateKeySpec = keyFactory.getKeySpec(privateKey, RSAPrivateKeySpec.class);
        RSAPublicKeySpec keySpec = new RSAPublicKeySpec(privateKeySpec.getModulus(), BigInteger.valueOf(65537));
        PublicKey publicKey = keyFactory.generatePublic(keySpec);
        //
        return new RSAItem(privateKey.getEncoded(), publicKey.getEncoded());
    }

//    public static void main(String[] args) throws NoSuchAlgorithmException {
//
//        String password = "";
//        password = RandomStringUtils.randomAlphanumeric(16);
//
//        KeyPair keyPair = generateKey(password);
//        byte[] pri = keyPair.getPrivate().getEncoded();
//        byte[] pub = keyPair.getPublic().getEncoded();
//
//        System.out.println("pwd=================");
//        System.out.println(password);
//        System.out.println("pub=================");
//        System.out.println(toBase64(pub));
//        System.out.println("pri=================");
//        System.out.println(toBase64(pri));
//    }

    @Getter
    public static final class RSAItem {
        private byte[] privateKey;
        private byte[] publicKey;

        public RSAItem(byte[] privateKey, byte[] publicKey) {
            this.privateKey = privateKey;
            this.publicKey = publicKey;
        }
    }

}
